package com.taoge.aop;

import com.taoge.annotation.SlideWindowLimit;
import com.taoge.exception.RateLimitException;
import com.taoge.properties.RateLimitProperties;
import org.springframework.aop.MethodBeforeAdvice;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.ZSetOperations;

import java.lang.reflect.Method;
import java.util.concurrent.TimeUnit;

/**
 * 滑动时间窗口限流切面
 * @author Mr.Tao
 */
public class SlideWindowLimitAspect implements MethodBeforeAdvice {

    private final RateLimitProperties limitProperties;

    private final RedisTemplate<String, Long> redisTemplate;

    public SlideWindowLimitAspect(RedisTemplate<String, Long> redisTemplate, RateLimitProperties limitProperties) {
        this.redisTemplate = redisTemplate;
        this.limitProperties = limitProperties;
    }

    @Override
    public void before(Method method, Object[] args, Object target) {
        // 获取LeakyBucketLimit注解，如果方法和类上同时出现，以方法上的为准
        Class<?> clazz = method.getDeclaringClass();
        SlideWindowLimit slideWindowLimit;
        slideWindowLimit = method.getAnnotation(SlideWindowLimit.class);
        if (slideWindowLimit == null) {
            slideWindowLimit = clazz.getAnnotation(SlideWindowLimit.class);
            if (slideWindowLimit == null) {
                return;
            }
        }

        // 根据用户IP（可选）和接口方法，构造key
        String key = limitProperties.getKey(slideWindowLimit.value(), slideWindowLimit.key(), clazz.getSimpleName(), method.getName(), limitProperties.getStrategyConfig().getSlideWindowLimitStrategy());

        // 限流逻辑实现
        ZSetOperations<String, Long> zSetOperations = redisTemplate.opsForZSet();
        // 记录本次访问的时间结点
        long currentTimeMillis = System.currentTimeMillis();
        zSetOperations.add(key, currentTimeMillis, currentTimeMillis);
        // 这一步是为了防止一直存在于内存中
        int time = slideWindowLimit.time();
        redisTemplate.expire(key, slideWindowLimit.time(), TimeUnit.SECONDS);
        // 移除{time}秒之前的访问记录（滑动窗口思想）
        zSetOperations.removeRangeByScore(key, 0, currentTimeMillis - time * 1000L);

        // 获得当前窗口内的访问记录数
        Long curCount = zSetOperations.zCard(key);
        // 限流判断
        if (curCount != null && curCount > slideWindowLimit.count()) {
            throw new RateLimitException("请求过于频繁，请稍后重试！");
        }
    }

}

